import numpy as np
from enum import Enum
import timeit
import time
import heapq
import matplotlib.pyplot as plt


class Node:
    def __init__(self, x, y, val):
        self.x = x
        self.y = y
        self.val = val
        self.fScore = np.inf
        self.gScore = np.inf
        self.key = str(x) + "," + str(y)

    def __lt__(self, rhs):
        """
        重写运算符。当两个Node的FScore相同时，需要按照Node大小进行排序，如果不重载运算符则无法比较两个Node的大小
        :param rhs:
        :return:
        """
        return self.val < rhs.val

    def setFScore(self, val):
        self.fScore = val

    def setGScore(self, val):
        self.gScore = val


class Rule(Enum):
    queen = 0  # 允许对角线
    rook = 1  # 只能上下左右


class AStar:
    def __init__(self, costArray, rule, p=0.1):
        self.costArray = costArray
        self.openSet = []
        self.closeSet = set()
        self.cameFrom = {}
        self.rule = rule
        self.D = np.quantile(costArray, p)
        self.D2 = self.D * np.sqrt(2)  # (1+np.sqrt(2))/2

    def find(self, startX, startY, endX, endY):
        """
        改编自：https://zhuanlan.zhihu.com/p/54510444
        * 初始化open_set和close_set；
        * 将起点加入open_set中，并设置gScore为0，计算hScore和fScore；
        * while open_set不为空
            * 从open_set中选取fScore最小的节点n：
            * 如果节点n为终点，则：
                * 从终点开始逐步追踪parent节点，一直达到起点；
                * 返回找到的结果路径，算法结束；
            * 如果节点n不是终点，则：
                * 将节点n从open_set中删除，并加入close_set中；
                * 遍历节点n所有的邻近节点：（在此处控制邻近规则）
                    * 如果邻近节点m在close_set中，则：
                        * 跳过，选取下一个邻近节点
                    * 如果邻近节点m不在open_set中，则：
                        * 设置节点m的parent为节点n
                        * 计算结点m的gScore：m.gScore = n.gScore + 从n到m的成本
                        * 计算节点m的fScore：m.fScore = m.gScore + 从m到终点的hScore
                        * 将节点m加入open_set中
                    * 如果邻近节点m在open_set中，则：
                        * 计算假设从结点n到结点m时结点m的tentative_gScore（用结点n对结点m进行松弛）：
                        * tentative_gScore = n.gScore + 从n到m的成本
                        * 如果tentative_gScore小于原来m的gScore，表明从结点n到结点m比原来到达结点m的实际路径短
                            * 将结点m的父节点设为n
                            * 更新结点m的gScore、fScore
        """
        self.openSet = []
        self.openSetMap = {}  # 行列号与节点的映射表，便于查找某个node是否在openset内
        self.closeSet = set()  # 哈希表，便于查找某个node是否在closeset内
        self.cameFrom = {}

        self.startNode = Node(startX, startY, self.costArray[startY, startX])
        self.goalNode = Node(endX, endY, self.costArray[endY, endX])
        self.startNode.setGScore(0)
        self.startNode.setFScore(self.computeHScore(self.startNode, self.goalNode))
        heapq.heappush(self.openSet, (self.startNode.fScore, self.startNode))
        self.openSetMap[self.startNode.key] = self.startNode  # 加入startnode的标志
        if self.rule == Rule.queen:
            return self.findQueen(endX, endY)
        else:
            return self.findRook(endX, endY)

    def findRook(self, endX, endY):
        while len(self.openSet) != 0:
            _, currentNode = heapq.heappop(self.openSet)  # 从最小堆里返回FScore最小的节点
            self.openSetMap.pop(currentNode.key)
            # print(currentNode.x, currentNode.y)
            if currentNode.x == endX and currentNode.y == endY:  # 是终点
                return currentNode.gScore, self.reconstructPath(currentNode)
            self.closeSet.add(currentNode.key)

            self.processNode(currentNode.x - 1, currentNode.y, currentNode)
            self.processNode(currentNode.x, currentNode.y - 1, currentNode)
            self.processNode(currentNode.x + 1, currentNode.y, currentNode)
            self.processNode(currentNode.x, currentNode.y + 1, currentNode)

        return -1, []

    def findQueen(self, endX, endY):
        while len(self.openSet) != 0:
            _, currentNode = heapq.heappop(self.openSet)  # 从最小堆里返回FScore最小的节点
            self.openSetMap.pop(currentNode.key)
            # print(currentNode.x, currentNode.y)
            if currentNode.x == endX and currentNode.y == endY:  # 是终点
                return currentNode.gScore, self.reconstructPath(currentNode)
            self.closeSet.add(currentNode.key)

            self.processNode(currentNode.x-1, currentNode.y+1, currentNode)
            self.processNode(currentNode.x-1, currentNode.y, currentNode)
            self.processNode(currentNode.x-1, currentNode.y-1, currentNode)
            self.processNode(currentNode.x, currentNode.y-1, currentNode)
            self.processNode(currentNode.x+1, currentNode.y-1, currentNode)
            self.processNode(currentNode.x+1, currentNode.y, currentNode)
            self.processNode(currentNode.x+1, currentNode.y+1, currentNode)
            self.processNode(currentNode.x, currentNode.y+1, currentNode)

        return -1, []

    def computeHScore(self, node, goal):
        if self.rule == Rule.rook:
            return self.D * (abs(node.x - goal.x) + abs(node.y - goal.y))
        elif self.rule == Rule.queen:
            diagonal = min(abs(node.x - goal.x), abs(node.y - goal.y))
            straight = abs(node.x - goal.x) + abs(node.y - goal.y)
            return self.D2 * diagonal + self.D * (straight - 2 * diagonal)

    def computeGScore(self, parent, child):
        if self.rule == Rule.rook:
            return (parent.val + child.val) / 2
        elif self.rule == Rule.queen:
            g = (parent.val + child.val) / 2
            if abs(parent.x - child.x) == 1 and abs(parent.y - child.y) == 1:
                g *= np.sqrt(2)
            return g

    def processNode(self, x, y, currentNode):
        if x < 0 or x > self.costArray.shape[1]-1 or y<0 or y>self.costArray.shape[0]-1:
            return
        if str(x) + "," + str(y) in self.closeSet:  # 在closeSet中
            return
        neighbor = self.openSetMap.get(str(x) + "," + str(y))
        if neighbor is None:  # 未访问过
            # print("未访问过：", x, y)
            neighbor = Node(x, y, self.costArray[y, x])
            self.cameFrom[neighbor] = currentNode
            neighbor.gScore = currentNode.gScore + self.computeGScore(currentNode, neighbor)
            neighbor.setFScore(neighbor.gScore + self.computeHScore(neighbor, self.goalNode))
            heapq.heappush(self.openSet, (neighbor.fScore, neighbor))
            self.openSetMap[neighbor.key] = neighbor
        else:  # 访问过
            tentative_gScore = currentNode.gScore + self.computeGScore(currentNode, neighbor)
            if tentative_gScore < neighbor.gScore:  # 从currentNode去neighbor比之前的路径好
                # print("更新最短路径：", x, y)
                self.cameFrom[neighbor] = currentNode
                neighbor.gScore = tentative_gScore
                # 修改二叉堆，删除旧的
                self.openSet.remove((neighbor.fScore, neighbor))
                neighbor.setFScore(neighbor.gScore + self.computeHScore(neighbor, self.goalNode))
                # 修改二叉堆，添加新的
                heapq.heappush(self.openSet, (neighbor.fScore, neighbor))

    def reconstructPath(self, currentNode):
        totalPath = [(currentNode.x, currentNode.y)]
        while currentNode in self.cameFrom.keys():
            currentNode = self.cameFrom[currentNode]
            totalPath.append((currentNode.x, currentNode.y))
        totalPath.reverse()
        return totalPath


if __name__ == "__main__":
    """
    # 读取tif数据：在QGIS的console里执行
    from osgeo import gdal
    import numpy as np
    ds = gdal.Open(r"D:\2021研一上\GIS开发\qgis-plugin\py\dem_tif\dem2.tif")
    myarray = np.array(ds.GetRasterBand(1).ReadAsArray())
    np.savetxt(r"D:\2021研一上\GIS开发\qgis-plugin\py\dem.txt", myarray,fmt="%.6f", delimiter=',')
    """
    # dem = np.loadtxt("data/dem_original.txt", delimiter=',')
    # astar = AStar(dem, Rule.rook)
    # number = 100
    # avg_time = timeit.timeit(lambda: astar.find(0, 100, 400, 300), number=number)/number  # 420, 60, 150, 160
    # # 画图
    # total_cost, path = astar.find(0, 100, 400, 300)
    # print("平均运行时间：", avg_time)
    # plt.imshow(dem)
    # plt.colorbar(fraction=0.03)
    # path = np.array(path)
    # plt.title(f"平均运行时间：{avg_time:.6f}s 总成本：{total_cost:.6f}")
    # plt.plot(path[:,0], path[:,1], c='r')
    # plt.show()
    # print("")

    dem = np.loadtxt("data/dem_original.txt", delimiter=',')
    costs = []
    times = []
    p_list = np.linspace(0, 0.5, 20)
    for p in p_list:
        astar = AStar(dem, Rule.rook, p)
        time1 = time.time()
        total_cost, path = astar.find(150,150,420,70)  # (0, 100, 400, 300)
        time2 = time.time()
        costs.append(total_cost)
        times.append(time2-time1)
    plt.plot(p_list, costs)
    plt.show()
    plt.plot(p_list, times)
    plt.show()
    pass
